/***********************************************
        File Name: cssp_server.cpp
        Author: Abby Cin
        Mail: abbytsing@gmail.com
        Created Time: 6/2/17 2:16 PM
***********************************************/

#include <boost/asio.hpp>
#include <iostream>
#include <mutex>
#include <queue>
#include <thread>
#include <sys/time.h>
#include <sys/resource.h>
#include "clp.h"
#include "logging.h"
#include "csssp_mysql.h"
#include "csssp_json.h"

using boost::asio::io_service;
using namespace boost::asio::ip;
using boost::system::error_code;
using boost::asio::async_read_until;
using boost::asio::async_write;
using namespace nm;

class Packet
{
  public:
    Packet()
      : header_size_(), data_(),  target_(), source_(),
        index_(), function_(), data_len_(), head_sum_()
    {
      for(auto&x: size_map_)
        header_size_ += x;
    }

    bool is_header_received()
    {
      return data_.size() >= header_size_;
    }

    bool is_header_valid()
    {
      unsigned checksum = head_checksum() % 0x100;
      return checksum == head_sum_;
    }

    bool is_data_valid()
    {
      int recv_sum = sum(data_.data(), data_len_) % 0x100;
      int data_sum = sum(data_.data() + data_len_, 1);
      return recv_sum == data_sum;
    }

    bool is_data_received()
    {
      return data_.size() >= data_len_ + 1;
    }

    std::string json_str()
    {
      return data_.substr(0, data_len_);
    }

    void clean_header()
    {
      data_.clear();
    }

    void clean_data()
    {
      data_.clear();
      target_ = 0;
      source_ = 0;
      index_ = 0;
      function_ = 0;
      data_len_ = 0;
      head_sum_ = 0;
    }

    void append(const char* data, size_t bytes)
    {
      data_.append(data, bytes);
    }

  private:
    /**
      FIXED_SIZE = 4,
      TARGET_SIZE = 2,
      SOURCE_SIZE = 2,
      INDEX_SIZE = 2,
      FUNC_SIZE = 1,
      DATA_SIZE = 2,
      HEADSUM_SIZE = 1,
    */
    const std::vector<int> size_map_{4, 2, 2, 2, 1, 2, 1};
    size_t header_size_;
    std::string data_;
    unsigned target_;
    unsigned source_;
    unsigned index_;
    unsigned function_;
    unsigned data_len_;
    unsigned head_sum_;

    unsigned head_checksum()
    {
      unsigned res = 0;
      auto iter = data_.data();
      size_t end = size_map_.size() - 1; // skip checksum filed
      for(size_t i = 0; i < end; ++i)
      {
        int tmp = sum(iter, size_map_[i]);
        res += tmp;
        iter += size_map_[i];
      }
      iter = data_.data();
      int i = 0; // skip fixed filed
      iter += size_map_[i++];
      target_ = uchar_to_int(iter, size_map_[i]);
      iter += size_map_[i++];
      source_ = uchar_to_int(iter, size_map_[i]);
      iter += size_map_[i++];
      index_ = uchar_to_int(iter, size_map_[i]);
      iter += size_map_[i++];
      function_ = uchar_to_int(iter, size_map_[i]);
      iter += size_map_[i++];
      data_len_ = uchar_to_int(iter, size_map_[i]);
      iter += size_map_[i++];
      head_sum_ = sum(iter, size_map_[i]);
      return res;
    }

    unsigned uchar_to_int(const char* data, int n)
    {
      unsigned res = 0;
      for(int i = 0; i < n; ++i)
      {
        res += (static_cast<unsigned char>(data[i]) << (i * 8));
      }
      return res;
    }

    unsigned sum(const char* data, int n)
    {
      unsigned res = 0;
      for(int i = 0; i < n; ++i)
      {
        res += static_cast<unsigned char>(data[i]);
      }
      return res;
    }
};

using Callback = std::function<void()>;

class Session : public std::enable_shared_from_this<Session>
{
  public:
    static std::shared_ptr<Session> make_session(io_service& io, int dur)
    {
      return std::make_shared<Session>(io, dur);
    }

    Session(io_service& io, int dur)
      : socket_(io), duration_(dur), time_point_(), cb_()
    {
      time_point_ = now();
    }

    ~Session()
    {
      if(socket_.is_open())
        socket_.close();
      log_info() << '[' << ip_ << ':' << port_ << "] connection closed";
    }

    tcp::socket& socket()
    {
      return socket_;
    }

    void set_callback(Callback&& cb)
    {
      cb_ = std::move(cb);
    }

    void run()
    {
      auto remote = socket_.remote_endpoint();
      ip_ = remote.address().to_string();
      port_ = remote.port();
      async_read();
    }

  private:
    enum { BUF_SIZE = 1024 * 10 };
    tcp::socket socket_;
    int duration_;
    std::chrono::steady_clock::time_point time_point_;
    Callback cb_;
    std::array<char, BUF_SIZE> buf_;
    Packet packet_;
    bool is_head_{true};
    std::string ip_{};
    int port_{0};

    auto now() -> decltype(std::chrono::steady_clock::now())
    {
      return std::chrono::steady_clock::now();
    }

    bool processing(size_t bytes)
    {
      for(size_t i = 0; i < bytes; ++i)
      {
        packet_.append(&buf_[i], 1);
        if(is_head_ && packet_.is_header_received())
        {
          if(packet_.is_header_valid())
          {
            packet_.clean_header();
            is_head_ = false;
          }
          else
          {
            log_err() << '[' << ip_ << ':' << port_ << "] invalid header";
            return false;
          }
        }
        if(!is_head_ && packet_.is_data_received())
        {
          if(packet_.is_data_valid())
          {
            is_head_ = true;
            auto json = packet_.json_str();
            if(cssp_data_parse(json.data(), json.size()) == -1)
            {
              log_err() << '[' << ip_ << ':' << port_ << "] insert to db failed";
              return false;
            }
            packet_.clean_data();
            break;
          }
          else
          {
            log_err() << '[' << ip_ << ':' << port_ << "] invalid data";
            return false;
          }
        }
      }
      return true;
    }

    void async_read()
    {
      auto self = shared_from_this();
      socket_.async_read_some(boost::asio::buffer(buf_),
        [this, self](const error_code& ec, size_t bytes) {
          if(ec)
          {
            if(ec != boost::asio::error::eof)
            {
              log_err() << '[' << ip_ << ':' << port_ <<"] error: " << ec.message();
            }
            return;
          }
          auto dur = std::chrono::duration_cast<std::chrono::seconds>(now() - time_point_);
          if(dur.count() >= duration_)
          {
            time_point_ = now();
            if(!processing(bytes))
              return;
          }
          async_read();
      });
    }
};

using SessionPtr = std::shared_ptr<Session>;

class Worker
{
public:
  Worker()
    : rbuf_(0),
    wbuf_(0),
    reader_(io_),
    writer_(io_),
    read_buf_(&rbuf_, sizeof rbuf_),
    write_buf_(&wbuf_, sizeof wbuf_),
    thread_([this] { thread_func(); })
    {
    }

  ~Worker()
  {
    if (thread_.joinable())
      thread_.join();
  }

  io_service& get_io_service()
  {
    return io_;
  }

  void enqueue(SessionPtr session)
  {
    enqueue_writer(session);
  }

  void stop()
  {
    std::lock_guard<std::mutex> lg(mtx_);
    io_.stop();
  }

private:
  char rbuf_, wbuf_;
  io_service io_;
  boost::asio::local::stream_protocol::socket reader_, writer_;
  boost::asio::mutable_buffers_1 read_buf_;
  boost::asio::const_buffers_1 write_buf_;
  std::thread thread_;
  std::queue<SessionPtr> tasks_;
  std::mutex mtx_;

  void enqueue_reader()
  {
    reader_.async_read_some(read_buf_,
      [this](const error_code& ec, size_t bytes)
    {
      read_handle(ec, bytes);
    });
  }

  void enqueue_writer(SessionPtr session)
  {
    std::lock_guard<std::mutex> lg(mtx_);
    tasks_.push(session);
    writer_.write_some(write_buf_);
  }

  void read_handle(const error_code& ec, size_t)
  {
    if (ec)
      return;
    std::lock_guard<std::mutex> lg(mtx_);
    while (!tasks_.empty())
    {
      auto session = tasks_.front();
      session->run();
      tasks_.pop();
    }
    enqueue_reader();
  }

  void thread_func()
  {
    boost::asio::local::connect_pair(reader_, writer_);
    enqueue_reader();
    io_.run();
  }
};

using WorkerPtr = std::shared_ptr<Worker>;

class Manager
{
  public:
    Manager(io_service& io, tcp::endpoint&& ep, unsigned num, int dur)
      : num_(num),
        duration_(dur),
        cur_(0),
        acceptor_(io, std::move(ep), true)
    {
      for(unsigned long i = 0; i < num_; ++i)
        workers_.emplace_back(std::make_shared<Worker>());

      acceptor_.listen();
      start_accept();
    }

    ~Manager()
    {
    }

    void stop()
    {
      acceptor_.cancel();
      for(auto& x: workers_)
        x->stop();
    }

  private:
    unsigned long num_;
    int duration_; // insert interval
    unsigned long cur_;
    tcp::acceptor acceptor_;
    std::vector<WorkerPtr> workers_;

    void accept_handle(SessionPtr session,
                       WorkerPtr worker,
                       const error_code& ec)
    {
      if(!ec)
      {
        ++cur_;
        worker->enqueue(session);
        start_accept();
      }
    }

    void start_accept()
    {
      if(cur_ == num_)
        cur_ = 0;
      auto worker = workers_.at(cur_);
      auto session(Session::make_session(worker->get_io_service(), duration_));
      acceptor_.async_accept(session->socket(),
                             [this, session, worker](const error_code& ec)
                             {
                               accept_handle(session, worker, ec);
                             });
    }
};

class Server
{
  public:
    Server(tcp::endpoint&& ep, unsigned num, int dur)
      : sg_(io_),
        manager_(io_, std::move(ep), num, dur)
    {
      sg_.add(SIGINT);
      sg_.add(SIGQUIT);
      sg_.add(SIGTERM);
      sg_.async_wait([this](const error_code& ec, int sig) {
        if(ec)
        {
          log_err() << ec.message();
        }
        log_info() << "caught signal " << sig << ", exit.";
        stop();
      });
    }

    void run()
    {
      io_.run();
    }

    void stop()
    {
      manager_.stop();
      io_.stop();
    }

  private:
    io_service io_;
    boost::asio::signal_set sg_;
    Manager manager_;
};

bool validate_ip(const nm::string_ext& ip)
{
  std::vector<nm::string_ext> v;
  ip.split(v, ".");
  if(v.size() != 4)
    return false;
  int count = 0;
  for(auto& x: v)
  {
    if(!x.is_digit())
      return false;
    if(x.size() > 3 || x.size() < 1)
      return false;
    count = 0;
    for(auto& c: x)
    {
      if(c > '9' || c < '0')
        return false;
      count *= 10;
      count += (c - '0');
    }
    if(count > 255)
      return false;
    // strict
    switch(x.size())
    {
      case 2:
        if(count < 10)
          return false;
        break;
      case 3:
        if(count < 100)
          return false;
        break;
      default:
        break;
    }
  }
  return true;
}

bool parse_cmd(Clp& cmd, std::string& csssp_host, int& csssp_port,
               std::string& db_host, int& db_port, std::string& db_user,
               std::string& db_name, std::string& db_passwd,
               int& thread_num, int& duration, unsigned long& nofile)
{
  if(auto res = cmd.get<std::string>("cssspd-host"))
  {
    csssp_host = res.value();
    if(!validate_ip(csssp_host))
    {
      std::cerr << "invalid -cssspd-host " << csssp_host << '\n';
      return false;
    }
  } else {
    return false;
  }

  if(auto res = cmd.get<int>("cssspd-port"))
  {
    csssp_port = res.value();
    if(csssp_port < 0 || csssp_port > 65535)
    {
      std::cerr << "invalid -cssspd-port " << csssp_port << '\n';
      return false;
    }
  } else {
    return false;
  }

  if(auto res = cmd.get<std::string>("db-host"))
  {
    db_host = res.value();
    if(!validate_ip(db_host))
    {
      std::cerr << "invalid -db-host " << db_host << '\n';
      return false;
    }
  } else {
    return false;
  }

  if(auto res = cmd.get<int>("db-port"))
  {
    db_port = res.value();
    if(db_port < 0 || db_port > 65535)
    {
      std::cerr << "invalid -db-port " << db_port << '\n';
      return false;
    }
  } else {
    return false;
  }

  if(auto res = cmd.get<std::string>("db-user"))
  {
    db_user = res.value();
  } else {
    return false;
  }

  if(auto res = cmd.get<std::string>("db-name"))
  {
    db_name = res.value();
  } else {
    return false;
  }

  if(auto res = cmd.get<std::string>("db-password"))
  {
    db_passwd = res.value();
  } else {
    return false;
  }

  if(auto res = cmd.get<int>("thread-num"))
  {
    thread_num = res.value();
    if(thread_num < 1)
    {
      std::cerr << "invalid -thread-num " << thread_num << '\n';
      return false;
    }
  } else {
    return false;
  }

  if(auto res = cmd.get<int>("insert-interval"))
  {
    duration = res.value();
    if(duration < 1)
    {
      std::cerr << "invalid -insert-interval " << duration << '\n';
      return false;
    }
  } else {
    return false;
  }

  if(auto res = cmd.get<unsigned long>("open-files"))
  {
    nofile = res.value();
    if(nofile < 1024)
    {
      std::cerr << "invalid -open-files " << nofile << '\n';
      return false;
    }
  } else {
    return false;
  }
  return true;
}


int loop(std::string& host, int port, int thread_num, int duration)
{
  try
  {
    auto endpoint = tcp::endpoint(address_v4::from_string(host), port);
    Server se(std::move(endpoint), thread_num, duration);
    se.run();
  }
  catch(const boost::system::system_error& e)
  {
    log_err() << e.what();
    std::cerr << e.what() << '\n';
    return 1;
  }
  return 0;
}

int try_lock_file(int fd)
{
  struct flock lock;
  memset(&lock, 0, sizeof(lock));

  lock.l_type = F_WRLCK;
  lock.l_whence = SEEK_SET;
  lock.l_start = 0;

  if(fcntl(fd, F_SETLK, &lock) == -1)
  {
    return -1;
  }
  return 0;
}

int already_running(const char* file, bool is_kill)
{
  errno = 0;
  int fd = open(file, O_CREAT | O_RDWR, S_IRUSR | S_IWUSR);
  if (fd == -1)
  {
    std::cerr << "open " << file << ' ' << strerror(errno) << ".\n";
    return -1;
  }
  char buf[10] = {0};
  if(is_kill)
  {
    if(try_lock_file(fd) == -1)
    {
      int pid = 0;
      if(read(fd, buf, sizeof(buf)) == -1)
      {
        std::cerr << "read " << strerror(errno) << ".\n";
        return -1;
      }
      sscanf(buf, "%d", &pid);
      if(pid != 0)
      {
        if(kill(pid, SIGTERM) == -1)
        {
          std::cerr << "kill pid " << pid << ' ' << strerror(errno);
          return -1;
        }
      }
    }
    else
    {
      std::cerr << "program is not running.\n";
    }
    close(fd);
    unlink(file);
  }
  else
  {
    if(try_lock_file(fd) == -1)
    {
      close(fd);
      std::cerr << "can't lock file! an instance is already running...\n";
      return -1;
    }
    sprintf(buf, "%d", getpid());

    // system will close fd when program exit and lock will be released automatically.
    if(write(fd, buf, strlen(buf)) == -1)
    {
      std::cerr << "write " << file  << ' ' << strerror(errno) << ".\n";
      return -1;
    }
  }
  return 0;
}

constexpr const char DEFAULT_PIDFILE_PATH[] = "/var/run/cssspd.pid";

int main(int argc, char* argv[])
{
  auto help = [](const char* cmd) {
    fprintf(stderr, "start: %s -c path_to_config\nstop: %s -stop\n", cmd, cmd);
  };
  if(argc < 2)
  {
    help(argv[0]);
    return 1;
  }
  Clp cli(argc, argv);
  cli.parse(Clp::SHORT);
  if(!cli.ok())
  {
    std::cerr << cli.msg() << '\n';
    return 1;
  }

  std::string cfg_path;
  if(cli.contain("stop") && argc == 2)
  {
    if(already_running(DEFAULT_PIDFILE_PATH, true) == -1)
      return 1;
    return 0;
  }
  else if(cli.get<std::string>("c") && argc == 3)
  {
    cfg_path = cli.get<std::string>("c").value();
  }
  else
  {
    help(argv[0]);
    return 1;
  }

  if(already_running(DEFAULT_PIDFILE_PATH, false) == -1)
    return 1;

  Clp cmd(cfg_path.c_str());
  cmd.parse(Clp::SHORT);
  if(!cmd.ok())
  {
    std::cerr << cmd.msg() << '\n';
    return 1;
  }

  std::string log_path{"/tmp"};
  std::string log_prefix{"cssspd"};
  long roll = 60 * 60 * 24; // rotate log file for every `roll` seconds

  if(auto res = cmd.get<std::string>("log-path"))
  {
    log_path = res.value();
  }
  if(auto res = cmd.get<std::string>("log-prefix"))
  {
    log_prefix = res.value();
  }
  if(auto res = cmd.get<long>("roll-interval"))
  {
    roll = res.value();
  }

  std::string csssp_host;
  int csssp_port;
  std::string db_host;
  int db_port;
  std::string db_user;
  std::string db_name;
  std::string db_passwd;
  int thread_num;
  int insert_interval;
  char error[512];
  unsigned long nofile;

  if(!parse_cmd(cmd, csssp_host, csssp_port, db_host, db_port, db_user,
                db_name, db_passwd, thread_num, insert_interval, nofile))
  {
    if(!cmd.ok()) {
      std::cerr << cmd.msg() << '\n';
    }
    return 1;
  }

  errno = 0;
  struct rlimit limit = {nofile, nofile};
  if(setrlimit(RLIMIT_NOFILE, &limit) == -1)
  {
    std::string msg{strerror(errno)};
    std::cerr << "setrlimit: " << cmd.msg() << '\n';
    return 1;
  }

  int ret = cssp_mysql_init(db_host.c_str(), db_user.c_str(),
                  db_passwd.c_str(), db_name.c_str(), db_port, error);
  if(ret < 0)
  {
    std::cerr << error << '\n';
    return 1;
  }

  if(!Logger<>::create_async(stderr, log_path.data(), log_prefix.data(), roll))
    return 1;

  return loop(csssp_host, csssp_port, thread_num, insert_interval);
}
